package main

import (
	"crypto/rand"
	"fmt"
	"github.com/ethereum/go-ethereum/common"
	gthCrypto "github.com/ethereum/go-ethereum/crypto"
	"github.com/protolambda/go-kzg/bls"
	"kzg-ipa-study/kzgX"
	"kzg-ipa-study/tools"
	"time"
)

func generalRoot(ins [][]byte) []bls.Fr {
	root := make([]bls.Fr, len(ins), len(ins))
	for i, one := range ins {
		hash := gthCrypto.Keccak256(one)
		kzgX.FrFromByte32(&root[i], hash)
	}
	return root
}

// zeroPolynomial returns the zero polynomial:
// z(x) = (x - z_0) (x - z_1) ... (x - z_{k-1})
func zeroPolynomial(zs []bls.Fr) []bls.Fr {
	var first bls.Fr
	bls.SubModFr(&first, &bls.ZERO, &zs[0])
	z := []bls.Fr{first, bls.ONE} // (x - z0)

	for i := 1; i < len(zs); i++ {
		var tmp bls.Fr
		bls.SubModFr(&tmp, &bls.ZERO, &zs[i])
		z = polynomialMul(z, []bls.Fr{tmp, bls.ONE}) // (x - zi)
	}
	return z
}

func polynomialMul(a, b []bls.Fr) []bls.Fr {
	r := arrayOfZeroes(len(a) + len(b) - 1)

	for i := 0; i < len(a); i++ {
		for j := 0; j < len(b); j++ {
			// r[i+j] = fAdd(r[i+j], fMul(a[i], b[j]))
			var tmp bls.Fr
			bls.MulModFr(&tmp, &a[i], &b[j])
			bls.AddModFr(&r[i+j], &r[i+j], &tmp)
		}
	}
	return r
}

func arrayOfZeroes(n int) []bls.Fr {
	r := make([]bls.Fr, n)
	for i := 0; i < n; i++ {
		r[i] = bls.ZERO
	}
	return r[:]
}

func testkzg() {
	s, err := tools.RandBigInt()
	if nil != err {
	}

	G1, G2 := kzgX.GenerateSetup(s.String(), 257)
	set := kzgX.NewKZGsetting(G1, G2, 4)

	data := make([][]byte, 256, 256)
	for i := 0; i < len(data); i++ { // 256
		seed := make([]byte, 32)
		rand.Read(seed)
		data[i] = seed
	}

	root := generalRoot(data)

	polynomial := zeroPolynomial(root)
	commitment := set.Commit(polynomial)
	timestart := time.Now()
	for i := 1; i < 1024; i++ { // 256*256*64 419万，40秒，约10万每秒
		polynomial = zeroPolynomial(root)
		commitment = set.Commit(polynomial)
	}
	fmt.Printf("%v\n", common.PrettyDuration(time.Since(timestart)))

	// P: 根据 trust setup 提供在 s 上的承诺（注： P 看不到s ）
	fmt.Printf("commitment: %s \n", bls.StrG1(commitment))

	// 以下是是证明生成和验证
	z := root[254]
	y := bls.ZERO

	// P：提供 f(z) = y 证据 witness
	proof := set.ComputeProofSingle(polynomial, z, y)

	// V: 对这次计算进行验证
	if !set.CheckProofSingle(commitment, proof, &z, &y) {
		fmt.Printf("承诺验证失败")
	}

	// 进行一次反例验证，seed 不在 root 中
	seed := make([]byte, 32)
	rand.Read(seed)
	var zz bls.Fr
	kzgX.FrFromByte32(&zz, seed)
	yy := bls.ZERO
	if set.CheckProofSingle(commitment, proof, &zz, &yy) {
		fmt.Printf("反例验证失败")
	}

}

func main() {
	testkzg()
}
